// Addin.cpp : Defines the exported functions for the DLL application.
//

#include "stdafx.h"
#include <RdcTransportLib/RdcTransport.h>
#include <plog/log.h>

#include <sstream>
#include <string>

//#include <boost/thread/thread.hpp>
#include <boost/thread.hpp>

//////////////////////////////////////////////////////////////////////////
namespace
{
    std::auto_ptr<rdc_transport::IVirtChannel>  channel;
    std::map<std::string, function_stub_t>      stubs;
    
    std::auto_ptr<boost::thread>                write_thread;
    std::auto_ptr<boost::thread>                write_thread1;

    boost::mutex                                on_read_mutex;


    void SendToVirtualChannel(std::vector<char>& buf)
    {
        channel->Write (&buf);
    }

    void create_stubs ()
    {
        stubs["accept"]        = &accept_stub;
        stubs["closesocket"]   = &closesocket_stub;
        stubs["bind"]          = &bind_stub;
        stubs["connect"]       = &connect_stub;
        stubs["freeaddrinfo"]  = &freeaddrinfo_stub;
        stubs["FreeAddrInfoW"] = &FreeAddrInfoW_stub;
        stubs["getaddrinfo"]   = &getaddrinfo_stub;
        stubs["GetAddrInfoW"]  = &GetAddrInfoW_stub;
        //?stubs["getnameinfo"]   = &getnameinfo_stub;
        //?stubs["GetNameInfoW"]  = &GetNameInfoW_stub;        
        stubs["getpeername"]   = &getpeername_stub;
        stubs["getsockname"]   = &getsockname_stub;
        stubs["ioctlsocket"]   = &ioctlsocket_stub;
        stubs["getsockopt"]    = &getsockopt_stub;
        stubs["listen"]        = &listen_stub;
        stubs["recv"]          = &recv_stub;
        stubs["recvfrom"]      = &recvfrom_stub;
        stubs["select"]        = &select_stub;
        stubs["send"]          = &send_stub;
        stubs["sendto"]        = &sendto_stub;
        stubs["shutdown"]      = &shutdown_stub;
        stubs["socket"]        = &socket_stub;
        stubs["setsockopt"]    = &setsockopt_stub;
        stubs["WSACleanup"]    = &WSACleanup_stub;
        //stubs["WSAIoctl"]      = &WSAIoctl_stub;
        //stubs["WSAJoinLeaf"]   = &WSAJoinLeaf_stub;
        stubs["__WSAFDIsSet"]    = &__WSAFDIsSet_stub;
        stubs["WSAGetLastError"] = &WSAGetLastError_stub;
        stubs["WSASetLastError"] = &WSASetLastError_stub;
        //stubs["WSARecv"]         = &WSARecv_stub;
        //stubs["WSARecvFrom"]     = &WSARecvFrom_stub;
        //stubs["WSASend"]         = &WSASend_stub;
        //stubs["WSASendTo"]       = &WSASendTo_stub;
        stubs["WSASocket"]      = &WSASocket_stub;
        stubs["WSAStartup"]      = &WSAStartup_stub;
        //?stubs["WSAGetOverlappedResult"] = &WSAGetOverlappedResult_stub;
    }

    struct Reader: public rdc_transport::IReadHandler
    {
        virtual void OnRead(std::vector<char>* pBuffer);
        virtual void OnReadFail(const char* what);
    } reader;

    void Reader::OnRead(rdc_transport::Buffer* pBuffer) throw()
    {
        boost::unique_lock<boost::mutex> lock(on_read_mutex);
        try
        {
            std::vector<char> buffer;
            buffer.swap(*pBuffer);

            Deserialize deserialize (buffer);
            
            LOG(plog::info) << "function : " << deserialize.get_header().func_name_ 
                            << "; TID : " << deserialize.get_header().TID_;

            std::vector<char> out_buffer;

            function_stub_t func = stubs[deserialize.get_header().func_name_];
            func (buffer, out_buffer);

            SendToVirtualChannel (out_buffer);
                            
            LOG(plog::info) << "OnRead";
        }
        catch(std::exception& ex)
        {
            LOG(plog::error) << ex.what();
        }
    }

    void Reader::OnReadFail(const char* what)
    {
        LOG(plog::error) << what;
    }

    struct Connection : public rdc_transport::IClientVirtChannelCallback
    {
        virtual void OnConnect() throw() 
        {
            create_stubs ();
        }

        virtual void OnDisconnect() throw() {}
        
        virtual void OnTerminate() throw() {}
    } connection;
}
//////////////////////////////////////////////////////////////////////////

BOOL VCAPITYPE VirtualChannelEntry(PCHANNEL_ENTRY_POINTS  pEntryPoints)
{
    try
    {
        channel =  rdc_transport::CreateClientChannel("WSCKDR", pEntryPoints, &connection, &reader, false);
        
        return TRUE;
    }
    catch(std::exception& ex)
    {
        LOG(plog::error) << ex.what();
        return false;
    }
}

std::wstring GetInstallDir ()
{
    const DWORD len = 256;
    wchar_t     env_var[len];
    
    int ret = ::GetEnvironmentVariableW (L"ProgramFiles", env_var, len);
    if (!ret && ERROR_ENVVAR_NOT_FOUND == GetLastError())
    {
        return std::wstring(L".");
    }
    return std::wstring(env_var);
}

//////////////////////////////////////////////////////////////////////////
BOOL APIENTRY DllMain( HMODULE hModule,
                       DWORD  ul_reason_for_call,
                       LPVOID lpReserved
					 )
{
	std::wstring log_file = GetInstallDir();
    log_file += L"\\Socket Redirection for RDS\\addin.csv";

    switch (ul_reason_for_call)
	{
	case DLL_PROCESS_ATTACH:
        plog::initialize_log(log_file.c_str(), plog::error);
        break;
	case DLL_PROCESS_DETACH:
        plog::release_log();
		break;
    case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
        break;
	}
	return TRUE;
}